import cv2
import numpy as np
import glob


# 加载标定相机的参数
with np.load("C.npz") as X:
    mtx,dist,_,_ = [X[i] for i in ("mtx",'dist', 'rvecs', 'tvecs')]

def draw(img,corners,imgpts):
    '''
    在图片上画出三维坐标轴
    :param img:原始图片
    :param corners:图像平面点坐标
    :param imgpts:三维点投影到二维图像平面点坐标
    :return:
    '''
    corner = tuple(corners[0].ravel())
    cv2.line(img,corner,tuple(imgpts[0].ravel()),(255,0,0),5)
    cv2.line(img,corner,tuple(imgpts[1].ravel()),(0,255,0),5)
    cv2.line(img,corner,tuple(imgpts[2].ravel()),(0,0,255),5)
    return img

# 棋盘格
w = 5
h = 5

# 初始化三维坐标点
objp = np.zeros((w*h,3),np.float32)
objp[:,:2] = np.mgrid[0:w,0:h].T.reshape(-1,2)

# 初始化三维坐标系
axis = np.float32([[3,0,0],[0,3,0],[0,0,-3]]).reshape(-1,3) #坐标轴

# 加载图片
path = glob.glob("E:\\VC_project\\figure\\chess2\\*.jpg")
i = 0
for fname in path:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    ret,corners = cv2.findChessboardCorners(gray,(w,h))

    if ret:
        # PnP 计算旋转向量和平移向量
        _, rvecs,tvecs,_ = cv2.solvePnPRansac(objp,corners,mtx,dist)
        print("*"*20)
        print("旋转向量",rvecs)
        print("平移向量",tvecs)

        # 计算三维点投影到二维图像平面的坐标
        imgpts,_ = cv2.projectPoints(axis,rvecs,tvecs,mtx,dist)
        # 将坐标显示在图片上
        img = draw(img,corners,imgpts)
        cv2.imwrite(str(i)+".jpg", img)
        i +=1
        cv2.imshow("img",img)
        cv2.waitKey(500)
cv2.waitKey()
